import halcon as ha
import os
from robot_arm import RobotArm
from alg import Alg
from threading import *
from vision_3d import V3D
from camera import Camera
import util_3d as u3d
import time as time

from spatialmath.base import *

def gen_pose_to_record_end(org_posture_rokea_type):
    pose_array = []
    base_t_end = u3d.rokae_pose_to_T(org_posture_rokea_type)
    # end_t_calib = u3d.xyzrpy_to_T(-0.175, 0.11, 0.08, 0.0, 0.0, 0.0)
    # end_t_calib = u3d.xyzrpy_to_T(-0.26, -0.13, 0.08, 0.0, 0.0, 0.0)
    end_t_calib = u3d.xyzrpy_to_T(0, 0, 0.0, 0.0, 0.0, 0.0)
    calib_t_end = trinv(end_t_calib)
    calib_aim_t_end_aim = calib_t_end
    for i in range(10):
        delta_calib = u3d.xyzrpy_to_T(0.0, 0.0, 0.0, 0.0, 0, -20+i*4)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(0.0, 0.0, 0.0, 0.0, -5+i*2, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)


    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(0.0, 0.0, 0.0, -5+i*2, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(-0.025+0.01*i, 0.0, 0.0, 0.0, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)


    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(0, -0.025+0.01*i, 0.0, 0.0, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(0, 0.0, -0.025+0.01*i, 0.0, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    return pose_array

def gen_pose_to_record_agv(org_posture_rokea_type):
    pose_array = []
    base_t_end = u3d.rokae_pose_to_T(org_posture_rokea_type)
    # end_t_calib = u3d.xyzrpy_to_T(-0.175, 0.11, 0.08, 0.0, 0.0, 0.0)
    end_t_calib = u3d.xyzrpy_to_T(0.26, -0.13, 0.08, 0.0, 0.0, 0.0)
    calib_t_end = trinv(end_t_calib)
    calib_aim_t_end_aim = calib_t_end
    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(0.0, 0.0, 0.0, 0.0, 0, -10+i*4)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(0.0, 0.0, 0.0, 0.0, -10+i*4, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(-0.025+0.01*i, 0.0, 0.0, 0.0, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    for i in range(5):
        delta_calib = u3d.xyzrpy_to_T(0, 0.0, -0.025+0.01*i, 0.0, 0, 0)
        calib_t_calib_aim = delta_calib
        base_t_end_aim = base_t_end @ end_t_calib @ calib_t_calib_aim @ calib_aim_t_end_aim
        base_t_end_aim_rokae = u3d.T_to_rokae_pose(base_t_end_aim)
        pose_array.append(base_t_end_aim_rokae)

    return pose_array
def gen_pose_to_record(org_posture_rokea_type , ):
    pose_array = []

    for i in range(20):
        org_posture_rokea_type_tmp = org_posture_rokea_type[:]
        if i == 0:
            org_posture_rokea_type_tmp[0] = org_posture_rokea_type_tmp[0] + 0.01
        if i == 1:
            org_posture_rokea_type_tmp[0] = org_posture_rokea_type_tmp[0] - 0.01
        if i == 2:
            org_posture_rokea_type_tmp[0] = org_posture_rokea_type_tmp[0] + 0.02
        if i == 3:
            org_posture_rokea_type_tmp[0] = org_posture_rokea_type_tmp[0] - 0.02
        if i == 4:
            org_posture_rokea_type_tmp[1] = org_posture_rokea_type_tmp[1] + 0.01
        if i == 5:
            org_posture_rokea_type_tmp[1] = org_posture_rokea_type_tmp[1] - 0.01
        if i == 6:
            org_posture_rokea_type_tmp[1] = org_posture_rokea_type_tmp[1] + 0.02
        if i == 7:
            org_posture_rokea_type_tmp[1] = org_posture_rokea_type_tmp[1] - 0.02
        if i == 8:
            org_posture_rokea_type_tmp[5] = org_posture_rokea_type_tmp[5] + 0.01
        if i == 9:
            org_posture_rokea_type_tmp[5] = org_posture_rokea_type_tmp[5] - 0.01
        if i == 10:
            org_posture_rokea_type_tmp[5] = org_posture_rokea_type_tmp[5] + 0.02
        if i == 11:
            org_posture_rokea_type_tmp[5] = org_posture_rokea_type_tmp[5] - 0.02

        if i == 12:
            org_posture_rokea_type_tmp[3] = org_posture_rokea_type_tmp[3] + 0.01
        if i == 13:
            org_posture_rokea_type_tmp[3] = org_posture_rokea_type_tmp[3] - 0.01
        if i == 14:
            org_posture_rokea_type_tmp[4] = org_posture_rokea_type_tmp[4] + 0.01
        if i == 15:
            org_posture_rokea_type_tmp[4] = org_posture_rokea_type_tmp[4] - 0.01

        if i == 16:
            org_posture_rokea_type_tmp[3] = org_posture_rokea_type_tmp[3] + 0.02
        if i == 17:
            org_posture_rokea_type_tmp[3] = org_posture_rokea_type_tmp[3] - 0.02
        if i == 18:
            org_posture_rokea_type_tmp[4] = org_posture_rokea_type_tmp[4] + 0.02
        if i == 19:
            org_posture_rokea_type_tmp[4] = org_posture_rokea_type_tmp[4] - 0.02
        pose_array.append(org_posture_rokea_type_tmp)

    return pose_array

# U3V:DA3797248 MV-CU060-10UM  AGV 上的
# U3V:DA3797129 MV-CU060-10UM  机械臂上的
def handeye_data_record(ca, arm:RobotArm, dir=''):
    ec = {}
    arm.robot.setDO(0, 0, True, ec)
    current_pose = arm.get_end_rokae_type()
    pose_array = gen_pose_to_record_end(current_pose)
    time_stamp = time.time()
    time_local = time.localtime(time_stamp)
    dir_name = ('./calib_data_' + str(time_local.tm_mon)+'_' +str(time_local.tm_mday) + '_' + str(time_local.tm_hour)
                + '_' + str(time_local.tm_min))
    os.mkdir(dir_name)

    for i in range(35):
        alpha = 57.295779513
        arm.movel_rokae(pose_array[i])
        arm.wait_util()
        time.sleep(1)
        pose_file_name = f"{i}.dat"
        image_file_name = f"{i}.png"
        pose_file_name = dir_name+'/'+ pose_file_name
        image_file_name = dir_name+'/'+ image_file_name
        posture = arm.get_end_rokae_type()
        img = ca.get_image()
        ha.write_image(img, 'png', 0, image_file_name)
        print(f"Record data  {i}")
        with open(pose_file_name, 'w') as f:
            text1 = f"r {alpha * posture[3]} {alpha * posture[4]} {alpha * posture[5]}\n"
            text2 = f"t {posture[0]} {posture[1]} {posture[2]}\n"
            f.write('f 2\n')
            f.write(text1)
            f.write(text2)
    print(f"Record data end")
    arm.movel_rokae(current_pose)
    arm.robot.setDO(0, 0, False, ec)

if __name__ == '__main__':

    # ca = Camera(camera_name='U3V:DA3797129 MV-CU060-10UM')
    # ca = Camera(camera_name='U3V:DA6258345 MV-CU060-10UM')
    ca = Camera(camera_name='U3V:DA5420519 MV-CU060-10UM')
    # ca.set_param()
    # ca.set_param_outer('param_outer_agv.txt')
    robot = RobotArm()

    t1 = Thread(target=ca.update_image_loop, daemon=True)
    t1.start()
    current_pose = robot.get_end_rokae_type()
    pose_array = gen_pose_to_record_end(current_pose)
    time_stamp = time.time()
    time_local = time.localtime(time_stamp)
    dir_name = ('./calib_data_' + str(time_local.tm_mon)+'_' +str(time_local.tm_mday) + '_' + str(time_local.tm_hour)
                + '_' + str(time_local.tm_min))
    os.mkdir(dir_name)

    for i in range(30):
        alpha = 57.295779513
        robot.movel_rokae(pose_array[i])
        robot.wait_util()
        time.sleep(1)
        pose_file_name = f"{i}.dat"
        image_file_name = f"{i}.png"
        pose_file_name = dir_name+'/'+ pose_file_name
        image_file_name = dir_name+'/'+ image_file_name
        posture = robot.get_end_rokae_type()
        img = ca.get_image()
        ha.write_image(img, 'png', 0, image_file_name)
        print(f"Record data  {i}")
        with open(pose_file_name, 'w') as f:
            text1 = f"r {alpha * posture[3]} {alpha * posture[4]} {alpha * posture[5]}\n"
            text2 = f"t {posture[0]} {posture[1]} {posture[2]}\n"
            f.write('f 2\n')
            f.write(text1)
            f.write(text2)
    print(f"Record data end")
    robot.movel_rokae(current_pose)








